package com.example.futureai.controller.chat;

import cn.hutool.core.map.MapUtil;
import cn.hutool.http.Header;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSON;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesis;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisOutput;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisParam;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisResult;
import com.alibaba.dashscope.utils.JsonUtils;
import com.example.futureai.constants.Result;
import com.example.futureai.dto.SendMsgDto;
import com.example.futureai.eum.ChatTYpe;
import com.example.futureai.service.CharRecordService;
import com.example.futureai.utils.TokenUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.dashscope.common.TaskStatus.SUCCEEDED;

/**
 * @ClassName: ImageAIGenerationController
 * @Description:
 * @Author: YourName
 * @Date: 2025/4/10
 */

@RestController
@Slf4j
@RequestMapping("/image")
public class ImageAIGenerationController {
    @Value("${aliblKey}")
    private String API_KEY;

    @Autowired
    private CharRecordService charRecordService;

    @PostMapping("/imageGeneration")
    public Result generateImage(@RequestBody SendMsgDto sendMsgDto) {
        sendMsgDto.setUserId(TokenUtil.getCurrentUserId());
        sendMsgDto.setChatType(ChatTYpe.IMAGE_CREATE.getType());
        //生成任务
        String taskId = this.createAsyncTask(sendMsgDto);
        //生成聊天记录
        sendMsgDto.setChatType(ChatTYpe.IMAGE_CREATE.getType());
        String mesgId = charRecordService.saveCharRecord(null, sendMsgDto);
        Map<String, String> map = new HashMap<>();
        map.put("taskId", taskId);
        map.put("mesgId", mesgId);
        return Result.OK(map);
    }

    @GetMapping("/taskStatus")
    public Result generateImage(String taskId, String mesgId) {
        return this.waitAsyncTask(taskId, mesgId);
    }

    /**
     * 创建异步任务
     *
     * @return taskId
     */
    public String createAsyncTask(SendMsgDto sendMsgDto) {
        String prompt = sendMsgDto.getInputInfo();
        //wanx2.1-t2i-turbo
        ImageSynthesisParam param =
                ImageSynthesisParam.builder()
                        .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                        .model("wanx2.1-t2i-plus")
                        .prompt(prompt)
                        .n(1)
                        .size("1024*1024")
                        .build();

        ImageSynthesis imageSynthesis = new ImageSynthesis();
        ImageSynthesisResult result = null;
        try {
            result = imageSynthesis.asyncCall(param);
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage());
        }
        System.out.println(JsonUtils.toJson(result));
        String taskId = result.getOutput().getTaskId();
        System.out.println("taskId=" + taskId);
        return taskId;
    }


    /**
     * 等待异步任务结束
     *
     * @param taskId 任务id
     */
    public Result waitAsyncTask(String taskId, String mesgId) {
        Map<String, String> res = new HashMap<>();
        String url = "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId;
        HttpResponse response = HttpRequest.get(url)
                .header(Header.AUTHORIZATION, "Bearer " + API_KEY)
                .execute();
        String imageUrl = "";
        if (response.getStatus() == 200) {
            JSONObject json = JSONUtil.parseObj(response.body());
            JSON outputJson = JSONUtil.parse(json.getStr("output"));
            ImageSynthesisOutput output = outputJson.toBean(ImageSynthesisOutput.class);
            String taskStatus = output.getTaskStatus();
            if (SUCCEEDED.getValue().equals(output.getTaskStatus())) {
                Map<String, String> map = output.getResults().get(0);
                imageUrl = MapUtil.getStr(map, "url");
                //将图片存储本地
                charRecordService.updateCharRecord(mesgId, imageUrl);
            }
            log.info("查询生成进度：taskStatus={}", taskStatus);
            res.put("taskId", taskId);
            res.put("taskStatus", taskStatus);
            res.put("imageUrl", imageUrl);
        }
        return Result.OK(res);
    }

}
